from rest_framework.views import APIView
from rest_framework.response import Response
from rest_framework.parsers import JSONParser
from django.conf import settings
import os
import subprocess
import re
import base64


class _getSm2KeyPair(APIView):
    def get(self, request):
        try:
            cmd_privatekey = "openssl genpkey -algorithm SM2 -out sm2_private_key.pem"
            cmd_publickey = "openssl pkey -pubout -in sm2_private_key.pem -out sm2_public_key.pem"
            subprocess.run(cmd_privatekey, shell=True, check=True)
            subprocess.run(cmd_publickey, shell=True, check=True)
            with open("sm2_private_key.pem", "r") as key_file:
                private_key = key_file.read()
                match = re.search(r"-----BEGIN PRIVATE KEY-----(.*?)-----END PRIVATE KEY-----", private_key, re.DOTALL)
                if match:
                    private_key = match.group(1).strip()
                else:
                    if os.path.exists("sm2_private_key.pem"):
                        os.remove("sm2_private_key.pem")
        
            with open("sm2_public_key.pem", "r") as key_file:
                public_key = key_file.read()
                match = re.search(r"-----BEGIN PUBLIC KEY-----(.*?)-----END PUBLIC KEY-----", public_key, re.DOTALL)
                if match:
                    public_key = match.group(1).strip()
                else:
                    if os.path.exists("sm2_private_key.pem"):
                        os.remove("sm2_private_key.pem")

            if os.path.exists("sm2_private_key.pem"):
                os.remove("sm2_private_key.pem")
            if os.path.exists("sm2_public_key.pem"):
                os.remove("sm2_public_key.pem")

            return Response({
                'PrivateKey': private_key,
                'PublicKey': public_key
            })
                
        except Exception as e:
            if os.path.exists("sm2_private_key.pem"):
                os.remove("sm2_private_key.pem")
            return Response({
                'error': str(e)
            })
        
class _sm2Encrypt(APIView):
    parser_classes = [JSONParser]
    def post(self, request):
        try:
            plaintext = request.data.get('plaintext', '').encode('utf-8')
            public_key = request.data.get('public_key', '').encode('utf-8')
            public_key_pem = f"-----BEGIN PUBLIC KEY-----\n{public_key.decode()}\n-----END PUBLIC KEY-----".encode('utf-8')
            with open("plaintext.txt", "wb") as file:
                file.write(plaintext)
            with open("sm2_public_key.pem", "wb") as file:
                file.write(public_key_pem)
            cmd = "openssl pkeyutl -encrypt -in plaintext.txt -pubin -inkey sm2_public_key.pem -out ciphertext.txt"
            subprocess.run(cmd, shell=True, check=True)
            
            with open("ciphertext.txt", "rb") as file:
                ciphertext = file.read()
                ciphertext = base64.b64encode(ciphertext).decode('utf-8')
            
            if os.path.exists("plaintext.txt"):
                os.remove("plaintext.txt")
            
            if os.path.exists("ciphertext.txt"):
                os.remove("ciphertext.txt")
            
            if os.path.exists("sm2_public_key.pem"):
                os.remove("sm2_public_key.pem")
            return Response({
                'ciphertext': ciphertext
            })
        
        except Exception as e:
            if os.path.exists("plaintext.txt"):
                os.remove("plaintext.txt")
            if os.path.exists("ciphertext.txt"):
                os.remove("ciphertext.txt")
            if os.path.exists("sm2_public_key.pem"):
                os.remove("sm2_public_key.pem")
            return Response({
                'error': str(e)
            })
class _sm2Decrypt(APIView):
    parser_classes = [JSONParser]
    def post(self, request):
        try:
            ciphertext = request.data.get('ciphertext', '').encode('utf-8')
            ciphertext = base64.b64decode(ciphertext)
            private_key = request.data.get('private_key', '').encode('utf-8')
            private_key_pem = f"-----BEGIN PRIVATE KEY-----\n{private_key.decode()}\n-----END PRIVATE KEY-----".encode('utf-8')
            with open("ciphertext.txt", "wb") as file:
                file.write(ciphertext)
            with open("sm2_private_key.pem", "wb") as file:
                file.write(private_key_pem)
            cmd = "openssl pkeyutl -decrypt -in ciphertext.txt -inkey sm2_private_key.pem -out plaintext.txt"
            subprocess.run(cmd, shell=True, check=True)
            
            with open("plaintext.txt", "rb") as file:
                plaintext = file.read()
                plaintext = plaintext.decode('utf-8')
            
            if os.path.exists("plaintext.txt"):
                os.remove("plaintext.txt")
            
            if os.path.exists("ciphertext.txt"):
                os.remove("ciphertext.txt")
            
            if os.path.exists("sm2_private_key.pem"):
                os.remove("sm2_private_key.pem")
            return Response({
                'plaintext': plaintext
            })
        
        except Exception as e:
            if os.path.exists("plaintext.txt"):
                os.remove("plaintext.txt")
            if os.path.exists("ciphertext.txt"):
                os.remove("ciphertext.txt")
            if os.path.exists("sm2_private_key.pem"):
                os.remove("sm2_private_key.pem")
            return Response({
                'error': str(e)
            })
class _sm2Sign(APIView):
    parser_classes = [JSONParser]
    def post(self, request):
        try:
            data = request.data.get('data', '').encode('utf-8')
            private_key = request.data.get('private_key', '').encode('utf-8')
            private_key_pem = f"-----BEGIN PRIVATE KEY-----\n{private_key.decode()}\n-----END PRIVATE KEY-----".encode('utf-8')

            with open("data.txt", "wb") as file:
                file.write(data)

            with open("sm2_private_key.pem", "wb") as file:
                file.write(private_key_pem)

            cmd = "openssl dgst -sm3 -sign sm2_private_key.pem -out signature.bin data.txt"
            subprocess.run(cmd, shell=True, check=True)

            with open("signature.bin", "rb") as file:
                signature = file.read()
                signature_base64 = base64.b64encode(signature).decode('utf-8')

            os.remove("data.txt")
            os.remove("sm2_private_key.pem")
            os.remove("signature.bin")

            return Response({
                'signature': signature_base64
            })

        except Exception as e:
            if os.path.exists("data.txt"):
                os.remove("data.txt")
            if os.path.exists("sm2_private_key.pem"):
                os.remove("sm2_private_key.pem")
            if os.path.exists("signature.bin"):
                os.remove("signature.bin")

            return Response({
                'error': str(e)
            })
class _sm2Verify(APIView):
    parser_classes = [JSONParser]

    def post(self, request):
        try:
            data = request.data.get('data', '').encode('utf-8')
            signature_base64 = request.data.get('signature', '')
            public_key = request.data.get('public_key', '').encode('utf-8')
            public_key_pem = f"-----BEGIN PUBLIC KEY-----\n{public_key.decode()}\n-----END PUBLIC KEY-----".encode('utf-8')

            with open("data.txt", "wb") as file:
                file.write(data)

            with open("sm2_public_key.pem", "wb") as file:
                file.write(public_key_pem)

            signature = base64.b64decode(signature_base64)
            with open("signature.bin", "wb") as file:
                file.write(signature)

            cmd = "openssl dgst -sm3 -verify sm2_public_key.pem -signature signature.bin data.txt"
            result = subprocess.run(cmd, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

            verified = "Verified OK" in result.stdout.decode()
            print(result.stdout.decode())

            os.remove("data.txt")
            os.remove("sm2_public_key.pem")
            os.remove("signature.bin")

            return Response({
                'verified': verified
            })

        except Exception as e:
            if os.path.exists("data.txt"):
                os.remove("data.txt")
            if os.path.exists("sm2_public_key.pem"):
                os.remove("sm2_public_key.pem")
            if os.path.exists("signature.bin"):
                os.remove("signature.bin")

            return Response({
                'error': str(e)
            })